import os
from dotenv import load_dotenv
from .ai_config import PROMPT_TEMPLATE
from .db_manager import DatabaseManager


# Load environment variables from .env file
load_dotenv()

class AIQueryBase:
    """
    公共AI交互基类，包含日志记录和SQL多轮验证等通用方法。
    """
    def __init__(self):
        """
        初始化创建数据库管理器
        """
        self.db = DatabaseManager()

    def _generate_query_prompt(self, question: str) -> str:
        """
        生成 AI 查询提示词，将表结构和问题拼接为完整 prompt。
        参数：
            question (str): 自然语言问题。
        返回：
            str: 用于 AI 的 prompt。
        """
        table_description = self.db.generate_table_description()
        return PROMPT_TEMPLATE.format(
            table_description=table_description,
            question=question
        )

    def log_ai_interaction(self, prompt: str, response_text: str, success: bool, error_msg: str = None):
        """
        记录每次AI请求的输入、输出、成功状态和错误信息到日志文件。
        日志文件超过10MB时自动新建带序号的新文件。
        """
        from datetime import datetime
        import os
        log_dir = "ai_logs"
        os.makedirs(log_dir, exist_ok=True)
        base_name = f"ai_interaction_{datetime.now().strftime('%Y%m%d')}"
        log_file = os.path.join(log_dir, f"{base_name}.log")
        max_size = 10 * 1024 * 1024  # 10MB
        file_index = 1
        # 检查文件大小，超过则新建带序号的文件
        while os.path.exists(log_file) and os.path.getsize(log_file) > max_size:
            log_file = os.path.join(log_dir, f"{base_name}_{file_index}.log")
            file_index += 1
        with open(log_file, 'a', encoding='utf-8') as f:
            f.write(f"时间: {datetime.now().isoformat()}\n")
            f.write(f"请求Prompt:\n{prompt}\n")
            f.write(f"AI返回: {'成功' if success else '失败'}\n")
            if response_text:
                f.write(f"输出内容:\n{response_text}\n")
            if error_msg:
                f.write(f"错误信息: {error_msg}\n")
            f.write("-" * 80 + "\n")
